import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
import numpy as np


# 1. 数据预处理
def getTrainSetAndTestSet(DataPath):
    data = pd.read_csv(DataPath)
    X = data[['AT', 'V', 'AP', 'RH']]  # 特征列
    y = data['PE']  # 标签列
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.25, random_state=1)
    print(f"训练集维度: {X_train.shape}, 测试集维度: {X_test.shape}")
    return X_train, X_test, y_train, y_test


# 2. 模型训练
def TrainLinearRegression(X_train, y_train):
    linreg = LinearRegression()
    linreg.fit(X_train, y_train)
    print(f"theta0: {linreg.intercept_:.4f}")
    print(f"theta1(AT): {linreg.coef_[0]:.4f}")
    print(f"theta2(V): {linreg.coef_[1]:.4f}")
    print(f"theta3(AP): {linreg.coef_[2]:.4f}")
    print(f"theta4(RH): {linreg.coef_[3]:.4f}")
    return linreg


# 3. 模型评估
def EvaluationModel(linreg, X_test, y_test):
    y_pred = linreg.predict(X_test)
    mse = np.mean((y_test - y_pred) ** 2)
    rmse = np.sqrt(mse)
    print(f"MSE: {mse:.4f}, RMSE: {rmse:.4f}")
    return y_pred


# 4. 可视化
def Visualization(y_test, y_pred):
    plt.figure(figsize=(8, 6))
    plt.scatter(y_test, y_pred, color='blue', alpha=0.5)
    plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=5)
    plt.xlabel("Measured")
    plt.ylabel("Predicted")
    plt.title("线性回归预测结果")
    plt.grid(True, alpha=0.3)
    plt.show()


# 主函数
if __name__ == "__main__":
    data_path = "Folds5x2_pp(1).csv"  # 数据集路径
    print("开始实验...")

    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(data_path)
    model = TrainLinearRegression(X_train, y_train)
    y_pred = EvaluationModel(model, X_test, y_test)
    Visualization(y_test, y_pred)

    print("实验完成")